BroadcastTo

将输入张量按 NumPy / MindSpore 广播规则扩展到目标形状 output_shape。 输入维数可少于输出维数(左侧隐式补 1);每一维上,输入长度必须为 1(可广播)或与输出长度相等。 形状维数不超过 8。

设输出维数为 \(M\),输入维数为 \(K\),且满足 \(K \le M\)。先将输入形状左侧补 1, 得到对齐后的输入形状 \(\tilde{I}\)

\[\begin{split}\tilde{I}_i = \begin{cases} 1, & 0 \le i < M - K \\ I_{i - (M - K)}, & M - K \le i < M \end{cases}\end{split}\]

对每个维度 \(i\),要求 \(\tilde{I}_i = 1\)\(\tilde{I}_i = O_i\), 其中 \(O\) 为输出形状。

索引映射

对输出线性下标 \(idx\),将其按输出 stride 分解为多维坐标 \((p_0,\ldots,p_{M-1})\), 再映射到输入偏移:

\[\begin{split}\begin{aligned} S^{out}_{M-1} &= 1,\quad S^{out}_{i} = S^{out}_{i+1} \cdot O_{i+1} \\ S^{in}_{M-1} &= \begin{cases} 0, & \tilde{I}_{M-1} = 1 \\ 1, & \text{otherwise} \end{cases},\quad S^{in}_{i} = \begin{cases} 0, & \tilde{I}_{i} = 1 \\ S^{in}_{i+1} \cdot \tilde{I}_{i+1}, & \text{otherwise} \end{cases} \\ \mathrm{in\_off} &= \sum_{i=0}^{M-1} p_i \cdot S^{in}_{i} \end{aligned}\end{split}\]

\(O[idx] = I[\mathrm{in\_off}]\)。当某维 \(\tilde{I}_i = 1\) 时, 该维 stride 为 0,对应元素沿该维重复。

输入:
  • input - 输入数据地址

  • input_shape - 输入形状数组

  • input_shape_size - 输入形状维数 \(K\)

  • output_shape - 目标输出形状数组

  • output_shape_size - 输出形状维数 \(M\)

  • core_mask - 核掩码(仅共享存储版本使用)

输出:
  • output - 输出数据地址

支持平台:

FT78NE MT7004

备注

  • FT78NE 支持 int8, int16, int32, fp32, fp64, cplx64, cplx128

  • MT7004 支持 int16, int32, fp16, fp32, cplx64

共享存储版本:

void i8_broadcast_to_s(int8_t *input, int8_t *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size, int core_mask)
void i16_broadcast_to_s(int16_t *input, int16_t *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size, int core_mask)
void i32_broadcast_to_s(int *input, int *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size, int core_mask)
void hp_broadcast_to_s(float16 *input, float16 *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size, int core_mask)
void fp_broadcast_to_s(float *input, float *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size, int core_mask)
void dp_broadcast_to_s(double *input, double *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size, int core_mask)
void c64_broadcast_to_s(float *input, float *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size, int core_mask)
void c128_broadcast_to_s(double *input, double *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size, int core_mask)

C调用示例:

 1// MT7004 示例(共享存储多核,DDR 地址)
 2// 例:input {1,N} -> output {N,N}
 3void TestBroadcastToSMCFp32(int n, int core_mask) {
 4    int core_id = get_core_id();
 5    int logic_core_id = GetLogicCoreId(core_mask, core_id);
 6    int core_num = GetCoreNum(core_mask);
 7    float *input = (float *)0x81000000;
 8    float *output = (float *)0x82000000;
 9    int input_shape[2] = {1, n};
10    int output_shape[2] = {n, n};
11    sys_bar(0, core_num);
12    fp_broadcast_to_s(input, output, input_shape, 2, output_shape, 2, core_mask);
13}
14
15void main() {
16    int core_mask = 0b1111;
17    TestBroadcastToSMCFp32(32, core_mask);
18}

私有存储版本:

void i8_broadcast_to_p(int8_t *input, int8_t *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size)
void i16_broadcast_to_p(int16_t *input, int16_t *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size)
void i32_broadcast_to_p(int *input, int *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size)
void hp_broadcast_to_p(float16 *input, float16 *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size)
void fp_broadcast_to_p(float *input, float *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size)
void dp_broadcast_to_p(double *input, double *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size)
void c64_broadcast_to_p(float *input, float *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size)
void c128_broadcast_to_p(double *input, double *output, int *input_shape, int input_shape_size, int *output_shape, int output_shape_size)

C调用示例:

 1// MT7004 示例(私有存储单核,AM 地址)
 2void TestBroadcastToAMFp32(int n) {
 3    float *input = (float *)0x10000000;
 4    float *output = (float *)0x10020000;
 5    int input_shape[2] = {1, n};
 6    int output_shape[2] = {n, n};
 7    fp_broadcast_to_p(input, output, input_shape, 2, output_shape, 2);
 8}
 9
10void main() {
11    TestBroadcastToAMFp32(32);
12}